package com.example.demo.ai.llm.service;

import com.jfinal.kit.Kv;
import com.jfinal.template.Engine;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;
import org.dromara.hutool.core.regex.ReUtil;
import org.dromara.hutool.core.text.StrUtil;
import org.noear.solon.ai.chat.ChatConfig;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.annotation.Component;
import org.noear.solon.flow.ChainContext;
import org.noear.solon.flow.Node;
import org.noear.solon.flow.TaskComponent;

/**
 * @author airhead
 */
@Component("intentionTask")
public class IntentionTask implements TaskComponent {
  private final Engine engine = Engine.create("llm");

  @Override
  public void run(ChainContext context, Node node) throws Throwable {
    Kv meta = Kv.create().set(node.meta());
    ChatConfig chatConfig = new ChatConfig();
    chatConfig.setModel(meta.getStr("model"));
    chatConfig.setProvider(meta.getStr("provider"));
    chatConfig.setApiUrl(meta.getStr("apiUrl"));
    chatConfig.setTimeout(Duration.ofSeconds(600));
    ChatModel chatModel = ChatModel.of(chatConfig).build();

    Kv model = Kv.create().set(context.model());

    List<ChatMessage> chatMessageList = new ArrayList<>();
    String systemTpl = meta.getStr("systemTpl");
    String systemContent =
        engine.getTemplateByString(systemTpl).renderToString(Kv.by("ctx", model.getStr("ctx")));
    ChatMessage system = ChatMessage.ofSystem(systemContent);
    chatMessageList.add(system);

    String inputKey = meta.getStr("input");
    String prompt = model.getStr(inputKey);
    String userTpl = meta.getStr("userTpl");
    String userContent =
        engine
            .getTemplateByString(userTpl)
            .renderToString(Kv.by("history", model.getStr("history")).set("prompt", prompt));

    ChatMessage userMessage = ChatMessage.ofUser(userContent);
    chatMessageList.add(userMessage);
    ChatResponse response = chatModel.prompt(chatMessageList).call();
    String content = response.getMessage().getContent();

    // 去掉think的部分。
    String pattern = "<think>.*?</think>";
    Pattern p = Pattern.compile(pattern, Pattern.DOTALL);
    content = StrUtil.trim(ReUtil.replaceAll(content, p, ""));

    String outputKey = meta.getStr("output");
    context.put(outputKey, content);

    Kv result = context.get("result");
    result.set("data", content);
  }
}
